
from fastapi import Body
from services.api_wrap import async_api_wrapper
from core.data_analyze.graph import build_graph
from fastapi.responses import StreamingResponse
import json

graph = build_graph()


@async_api_wrapper
async def generate_report(
    task_description: str = Body('', description='任务描述'),
    datasets: list = Body([], description='任务相关的数据列表')
):
    response = await graph.ainvoke(
        input={
            "auto_decision": True,
            "max_retry_times": 3,
            "cur_retry_times": 0,
            "task_description": task_description,
            "datasets": datasets
        },
        stream_mode=["custom"])
    return response[-1][1]

@async_api_wrapper
async def stream_generate_report(
    task_description: str = Body('', description='任务描述'),
    datasets: list = Body([], description='任务相关的数据列表')
):
    import time
    def data_stream():
        for event in graph.stream(
            input={
                'task_description': task_description,
                'datasets': datasets
            },
            stream_mode=["custom"]
        ):
            print(event)
            time.sleep(1)
            yield json.dumps(event, ensure_ascii=False).encode('utf-8')

    return StreamingResponse(data_stream(), media_type="text/plain")


@async_api_wrapper
async def request_handler(
        query: str = Body(..., description='任意的文本'),
        times: int = Body(default=10, description='打印次数'),
):
    import time
    def ask_statesman():
        for k in range(times):
            event = {k: query}
            print(event)
            time.sleep(1)
            yield json.dumps(event, ensure_ascii=False)

    return StreamingResponse(ask_statesman(), media_type="text/plain")
